#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据导入脚本 - 将 Excel 文件导入到 SQLite 数据库

功能：
1. 读取 Data Source 文件夹中的 Sample.xlsx 和基础表.xlsx
2. 将数据导入到 SQLite 数据库中
3. 自动创建表结构
4. 处理中文列名和数据
"""

import os
import sqlite3
import pandas as pd
from pathlib import Path
import logging
from typing import Dict, List, Any

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('import_log.txt', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class ExcelToSQLiteImporter:
    """Excel 到 SQLite 数据导入器"""
    
    def __init__(self, db_path: str = "Chinook.db"):
        """初始化导入器
        
        Args:
            db_path: SQLite 数据库文件路径
        """
        self.db_path = db_path
        self.project_root = Path(__file__).parent.parent
        self.data_source_dir = self.project_root / "Data Source"
        
    def clean_column_name(self, col_name: str) -> str:
        """清理列名，使其适合作为 SQL 列名
        
        Args:
            col_name: 原始列名
            
        Returns:
            清理后的列名
        """
        # 移除特殊字符，替换为下划线
        import re
        cleaned = re.sub(r'[^\w\u4e00-\u9fff]', '_', str(col_name))
        # 移除开头和结尾的下划线
        cleaned = cleaned.strip('_')
        # 如果列名为空或只有下划线，使用默认名称
        if not cleaned or cleaned == '_':
            cleaned = 'column_' + str(hash(col_name) % 1000)
        return cleaned
    
    def get_sql_type(self, dtype: str) -> str:
        """根据 pandas 数据类型获取对应的 SQL 类型
        
        Args:
            dtype: pandas 数据类型
            
        Returns:
            SQL 数据类型
        """
        if 'int' in str(dtype):
            return 'INTEGER'
        elif 'float' in str(dtype):
            return 'REAL'
        elif 'datetime' in str(dtype):
            return 'DATETIME'
        elif 'bool' in str(dtype):
            return 'BOOLEAN'
        else:
            return 'TEXT'
    
    def create_table_sql(self, table_name: str, df: pd.DataFrame) -> str:
        """生成创建表的 SQL 语句
        
        Args:
            table_name: 表名
            df: DataFrame
            
        Returns:
            CREATE TABLE SQL 语句
        """
        columns = []
        for col in df.columns:
            clean_col = self.clean_column_name(col)
            sql_type = self.get_sql_type(df[col].dtype)
            columns.append(f'`{clean_col}` {sql_type}')
        
        columns_sql = ',\n    '.join(columns)
        return f"""CREATE TABLE IF NOT EXISTS `{table_name}` (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    {columns_sql}
)"""
    
    def import_excel_file(self, file_path: Path, table_name: str) -> bool:
        """导入单个 Excel 文件
        
        Args:
            file_path: Excel 文件路径
            table_name: 目标表名
            
        Returns:
            是否导入成功
        """
        try:
            logger.info(f"开始导入文件: {file_path}")
            
            # 读取 Excel 文件
            # 尝试读取第一个工作表
            df = pd.read_excel(file_path, sheet_name=0)
            
            # 检查是否有数据
            if df.empty:
                logger.warning(f"文件 {file_path} 没有数据")
                return False
            
            # 清理列名
            df.columns = [self.clean_column_name(col) for col in df.columns]
            
            # 处理 NaN 值
            df = df.fillna('')
            
            logger.info(f"文件 {file_path} 包含 {len(df)} 行数据，{len(df.columns)} 列")
            logger.info(f"列名: {list(df.columns)}")
            
            # 连接数据库
            with sqlite3.connect(self.db_path) as conn:
                # 创建表
                create_sql = self.create_table_sql(table_name, df)
                logger.info(f"创建表 SQL: {create_sql}")
                conn.execute(create_sql)
                
                # 清空表（如果需要重新导入）
                conn.execute(f"DELETE FROM `{table_name}`")
                
                # 导入数据
                df.to_sql(table_name, conn, if_exists='append', index=False)
                
                # 验证导入结果
                cursor = conn.execute(f"SELECT COUNT(*) FROM `{table_name}`")
                count = cursor.fetchone()[0]
                logger.info(f"成功导入 {count} 行数据到表 {table_name}")
            
            return True
            
        except Exception as e:
            logger.error(f"导入文件 {file_path} 失败: {str(e)}")
            return False
    
    def import_all_files(self) -> Dict[str, bool]:
        """导入所有指定的 Excel 文件
        
        Returns:
            导入结果字典
        """
        results = {}
        
        # 定义要导入的文件和对应的表名
        files_to_import = [
            ("Sample.xlsx", "sample_data"),
            ("基础表.xlsx", "base_data")
        ]
        
        for filename, table_name in files_to_import:
            file_path = self.data_source_dir / filename
            
            if not file_path.exists():
                logger.error(f"文件不存在: {file_path}")
                results[filename] = False
                continue
            
            results[filename] = self.import_excel_file(file_path, table_name)
        
        return results
    
    def show_database_info(self):
        """显示数据库信息"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                # 获取所有表
                cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
                tables = cursor.fetchall()
                
                logger.info("=== 数据库信息 ===")
                logger.info(f"数据库文件: {self.db_path}")
                logger.info(f"表数量: {len(tables)}")
                
                for table in tables:
                    table_name = table[0]
                    # 获取表结构
                    cursor = conn.execute(f"PRAGMA table_info(`{table_name}`)")
                    columns = cursor.fetchall()
                    
                    # 获取行数
                    cursor = conn.execute(f"SELECT COUNT(*) FROM `{table_name}`")
                    row_count = cursor.fetchone()[0]
                    
                    logger.info(f"\n表名: {table_name}")
                    logger.info(f"行数: {row_count}")
                    logger.info(f"列数: {len(columns)}")
                    logger.info("列信息:")
                    for col in columns:
                        logger.info(f"  - {col[1]} ({col[2]})")
                    
                    # 显示前几行数据示例
                    if row_count > 0:
                        cursor = conn.execute(f"SELECT * FROM `{table_name}` LIMIT 3")
                        sample_data = cursor.fetchall()
                        logger.info("数据示例:")
                        for i, row in enumerate(sample_data, 1):
                            logger.info(f"  行{i}: {row[:5]}...")  # 只显示前5列
                
        except Exception as e:
            logger.error(f"显示数据库信息失败: {str(e)}")

def main():
    """主函数"""
    print("=== Excel 到 SQLite 数据导入工具 ===")
    print("正在导入 Sample.xlsx 和基础表.xlsx 到 SQLite 数据库...\n")
    
    # 创建导入器
    importer = ExcelToSQLiteImporter("Chinook.db")
    
    # 导入所有文件
    results = importer.import_all_files()
    
    # 显示结果
    print("\n=== 导入结果 ===")
    for filename, success in results.items():
        status = "成功" if success else "失败"
        print(f"{filename}: {status}")
    
    # 显示数据库信息
    print("\n=== 数据库信息 ===")
    importer.show_database_info()
    
    print(f"\n数据库文件已保存为: {importer.db_path}")
    print("导入完成！")

if __name__ == "__main__":
    main()